import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt

def historical_performance(code_list, start_date, end_date):
    """
    获取股票的历史表现
    :param code_list: 股票代码列表
    :param start_date: 开始日期
    :param end_date: 结束日期
    :return: 股票的历史表现
    """
    # 拉取k线数据
    k_data_dict = {}
    print('获取k线数据')
    index = 1
    for code in code_list:
        print(f'{index}. {code}')
        index += 1
        rs = bs.query_history_k_data_plus(code,
                            "date,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                            start_date=start_date, end_date=end_date,
                            frequency="d", adjustflag="3")
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)
        k_data_dict[code] = {}
        k_data_dict[code]['k_data'] = result
    # 根据期初股价进行配置
    open_sum = 0
    for code in code_list:
        open_sum += float(k_data_dict[code]['k_data']['open'][0])
    for code in code_list:
        k_data_dict[code]['weight'] = open_sum / float(k_data_dict[code]['k_data']['open'][0])
    print(k_data_dict)
    # 从日期的维度聚合总价值
    date_value_data = {}
    for code in code_list:
        code_k_data = k_data_dict[code]
        weight = k_data_dict[code]['weight']
        for index, row in code_k_data['k_data'].iterrows():
            date = row['date']
            if date in date_value_data:
                date_value_data[date] += float(row['close']) * weight
            else:
                date_value_data[date] = float(row['close']) * weight
    # 转换为list并按date升序排序
    date_value_list = []
    for date, value in date_value_data.items():
        date_value_list.append({'date': date, 'value': value})
    date_value_list.sort(key=lambda x: x['date'])
    # 转换为收益率
    profit_rate = to_rate(date_value_list)
    # 转换为DataFrame
    df = pd.DataFrame(profit_rate)
    # 打印输出最终的收益率
    print(f'策略收益率为：{df.iloc[-1]["rate"] * 100.0:.2f}%')
    # 获取沪深300数据
    hs300_data = bs.query_history_k_data_plus("sh.000300",
            "date,code,open,high,low,close,preclose,volume,amount,pctChg",
            start_date=start_date, end_date=end_date, frequency="d")
    data_list = []
    while (hs300_data.error_code == '0') & hs300_data.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(hs300_data.get_row_data())
    result = pd.DataFrame(data_list, columns=hs300_data.fields)
    # 转换为收益率
    hs300_date_list = result[['date', 'close']].values.tolist()
    hs300_date_list = list(map(lambda r : {'date': r[0], 'value': float(r[1])}, hs300_date_list))
    hs300_rate = pd.DataFrame(to_rate(hs300_date_list))
    print(f'沪深300收益率为：{hs300_rate.iloc[-1]["rate"] * 100.0:.2f}%')
    # 绘制折线图
    plt.plot(df['date'], df['rate'], label='Strategy 1')
    plt.plot(hs300_rate['date'], hs300_rate['rate'], label='HS300')
    plt.xlabel('Date')
    plt.ylabel('Profit Rate(%)')
    plt.legend()
    plt.title(f'Historical Performance {start_date} - {end_date}')
    plt.show()

def to_rate(data):
    # 转换为收益率
    profit_rate = [{'date': data[0]['date'], 'rate': 0}]
    initial_amount = data[0]['value']
    for i in range(1, len(data)):
        rate = (data[i]['value'] - initial_amount) / initial_amount
        profit_rate.append({'date': data[i]['date'], 'rate': rate})
    return profit_rate

codes = pd.read_csv('~/Desktop/low_pr.csv')['code'].to_list()

bs.login()
historical_performance(codes, '2021-01-01', '2025-08-25')
bs.logout()